#include "Controller.h"

#include <iostream>
#include <vector>
#include <sstream>
#include <cstring>
#include <iostream>
#include <fstream>
#include "Config.h"
#include "io.h"
#include "WordsLine.h"
#include "PamImage.h"
#include "Test.h"
#include "Database.h"
#include "Segmentator.h"
#include "Preprocessor.h"

#define DEBUG
string databasePrefix = "database";

void trainNFold(string imagesPath, string inputWordsPath, int n){
  //      extract features for each word
  //      write database
  cout << endl << "//// I'm going to train a databases.////" << endl << endl;
  
  vector< string > imageFilenames = readDirectory(imagesPath);
  vector< string > wordFilenames = readDirectory(inputWordsPath);
 
  vector<Database> databases;
  for(int nIndex = 0; nIndex < n; nIndex++){
    Database database;
    databases.push_back(database);
  }
  
  for(size_t fileNr = 0; fileNr < imageFilenames.size(); fileNr++)
  {
    cout << "processing file " << wordFilenames[fileNr] << endl;
    PamImage fullColorImage( imageFilenames[fileNr] );   
    
    cout << "reading xml" << wordFilenames[fileNr] << endl;
    vector<WordsLine> lineInfo = readXml( wordFilenames[fileNr] ); 
    
    cout << "extracting line images " << imageFilenames[fileNr] << endl;
    ExtractLineImagesFromDocument(&fullColorImage, &lineInfo);
    cout << "extracting word images " << imageFilenames[fileNr] << endl;
    ExtractWordImages(&lineInfo);

    
    for(size_t lineNr = 0; lineNr < lineInfo.size(); lineNr++)
    {           
      vector<Word> wordInfo = lineInfo[lineNr].words;      
      for(size_t wordNr = 0; wordNr < wordInfo.size(); wordNr++)
      {
	//process each word
	string target = wordInfo[wordNr].text;	
	//extract features
	Descriptors wordDescr = getDescriptors( wordInfo[wordNr].image );
	//add to database    
	for(int nIndex = 0; nIndex < n; nIndex++){
	  if(isInNthTestSet(n, nIndex, fileNr)){
	    continue;	    
	  }
	  databases[nIndex].addPrototype(target, wordDescr);
	}	
      }      
    }
  }
  //write to file
  #ifdef DEBUG
  cout << "save databases" << endl;
  #endif
  for(int nIndex = 0; nIndex < n; nIndex++){
    databases[nIndex].saveDatabase(getNFoldDatabaseName(n,nIndex));
  }
}

bool isInNthTestSet(int n, int nIndex, int index){
  return index % n == nIndex;
}

string getNFoldDatabaseName(int n, int number){
    stringstream databaseFileName;
    databaseFileName << databasePrefix 
      << "_" << Config::getInstance()->radonanglecount
      << "_" << Config::getInstance()->radonlinecount
      << "_" << Config::getInstance()->radonhistowidth
      << "_" << Config::getInstance()->lowResolutionWordWidth
      << "_" << Config::getInstance()->lowResolutionWordHeight
      << "_" << Config::getInstance()->bottomProfileSampleNumber
      << "_" << n << "fold_number"<<number<<".db";
    return databaseFileName.str();
}

string getNFoldOutputFolderName(int n, int number){
    stringstream outputFolderName;
    outputFolderName << "output_" << n << "fold_number"<<number;
    return outputFolderName.str();
}

void train(string imagesPath, string inputWordsPath, string outputDatabaseName)
{
  //      extract features for each word
  //      write database
  cout << endl << "//// I'm going to train a database. Will write to " << outputDatabaseName << " ////" << endl << endl;
  
  vector< string > imageFilenames = readDirectory(imagesPath);
  vector< string > wordFilenames = readDirectory(inputWordsPath);
  
  Database database;
  
  for(size_t fileNr = 0; fileNr < imageFilenames.size(); fileNr++)
  {
    cout << "processing file " << wordFilenames[fileNr] << endl;
    PamImage fullColorImage( imageFilenames[fileNr] );   
    
    cout << "reading xml" << wordFilenames[fileNr] << endl;
    vector<WordsLine> lineInfo = readXml( wordFilenames[fileNr] ); 
    
    cout << "extracting line images " << imageFilenames[fileNr] << endl;
    ExtractLineImagesFromDocument(&fullColorImage, &lineInfo);
    cout << "extracting word images " << imageFilenames[fileNr] << endl;
    ExtractWordImages(&lineInfo);



    #ifdef DEBUG
    cout << "save images" << endl;    
    for(size_t lineNr = 0; lineNr < lineInfo.size(); lineNr++)
    {       
      vector<Word> wordInfo = lineInfo[lineNr].words;      
      for(size_t wordNr = 0; wordNr < wordInfo.size(); wordNr++)
      {    
	      SaveWordImage("trainwordimages", wordInfo[wordNr], wordInfo[wordNr].image);
      }
      SaveLineImage("trainlineimages", lineInfo[lineNr], lineInfo[lineNr].image);
    }
    cout << "add descriptors" << endl;
    #endif
    
    for(size_t lineNr = 0; lineNr < lineInfo.size(); lineNr++)
    {           
      vector<Word> wordInfo = lineInfo[lineNr].words;      
      for(size_t wordNr = 0; wordNr < wordInfo.size(); wordNr++)
      {
      	//process each word
	string target = wordInfo[wordNr].text;	
        //extract features
	Descriptors wordDescr = getDescriptors( wordInfo[wordNr].image );
	//add to database    
      	database.addPrototype(target, wordDescr);
        
      }
      
    }
  }
  //write to file
  #ifdef DEBUG
  cout << "save database" << endl;
  #endif
  database.saveDatabase(outputDatabaseName);

 // Database testdatabase;
 // testdatabase.addDatabaseFile(outputDatabaseName);
 // testdatabase.printData();
}

void recognize(string databaseName, string imageName, string inputLineName, string outputWordsName)
{
    //TODO: read database
    //      read image and line-info --> classify to vector<WordsLine> using database
    //      write xmlFile
    cout << endl << "//// I'm going to classify " << imageName << " ////" << endl << endl;
    
    cout << "Reading database " << databaseName << " .... " << endl;
    Database database(databaseName);
    
    cout << "Reading image " << imageName << " .... " << endl;
    PamImage image(imageName);  
    
    cout << "Reading words " << inputLineName << " .... " << endl;
    vector<WordsLine> lineInfo = readXml(inputLineName);
    
    ExtractLineImagesFromDocument(&image, &lineInfo); //extracts the line images from the image
    DetermineWordPositions(&lineInfo); //determines the word hypotheses
    ExtractWordImages(&lineInfo);//extracts the images with the coordinates determined in the previous step
    
    cout << "Classifying words ..." << endl;
    for(size_t lineNr = 0; lineNr < lineInfo.size(); lineNr++)
    {               
      for(size_t wordNr = 0; wordNr < lineInfo[lineNr].words.size(); wordNr++)
      {
	Descriptors wordDescr = getDescriptors( lineInfo[lineNr].words[wordNr].image );
	//SaveWordImage("LowResolution", lineInfo[lineNr].words[wordNr], *lowResolutionToImage(wordDescr.lowResolution));
	vector<Match> matches = database.getBestMatches(wordDescr);
	if(matches.size() ==0){
	  cout << "no match found" << endl;
	  continue;
	}
	string text = matches[0].text;
	double score = matches[0].score;
	cout << " word classified as \"" << text << "\"" << " with score " << score << " (lower is better)" << endl;
	lineInfo[lineNr].words[wordNr].text = text;
	lineInfo[lineNr].words[wordNr].score = score;
      }
    
      #ifdef DEBUG
      //std::ostringstream lineNrStr; lineNrStr << lineNr;
      cout << "save images" << endl;
      cout.flush();
      vector<Word> wordInfo = lineInfo[lineNr].words;      
      for(size_t wordNr = 0; wordNr < wordInfo.size(); wordNr++){
	SaveWordImage("overlappingwordimages", wordInfo[wordNr], wordInfo[wordNr].image);
      }
      #endif
      
      cout << "construct non overlapping line with smallest score" << endl;
      cout.flush();
     // lineInfo[lineNr].words = createLineWithLowestScore(lineInfo[lineNr].words);
      lineInfo[lineNr].words = GetBestLine(lineInfo[lineNr].words);
      
      cout << endl << endl << "BEST LINE:" << endl;
      for(unsigned i = 0; i < lineInfo[lineNr].words.size(); i++){
	cout << lineInfo[lineNr].words[i].text << "  ";
      }
      cout << endl;
      #ifdef DEBUG
      cout << "save images" << endl;
      cout.flush();
      wordInfo = lineInfo[lineNr].words;      
      for(size_t wordNr = 0; wordNr < wordInfo.size(); wordNr++){
	SaveWordImage("recognizewordimages", wordInfo[wordNr], wordInfo[wordNr].image);
	cout.flush();
      }
      SaveLineImage("recognizelineimages", lineInfo[lineNr], lineInfo[lineNr].image);
      cout.flush();
      #endif
    }    
    cout << "save words file" << endl;
    cout.flush();
    writeXml(lineInfo, outputWordsName);
}

void test(string realWordsPath, string ourWordsPath, bool matchcontent)
{
  //TODO: read in files from both directories
  //      compare the output of our program with the correct output
  
  cout << endl << "//// I'm going to test " << realWordsPath << " with " << ourWordsPath << " ////" << endl << endl;
  if(matchcontent){
    cout << "Matching content" << endl << endl;
  }else{
    cout << "Not matching content" << endl << endl;
  }

  vector< string > ourFilenames = readDirectory(ourWordsPath);
  vector< string > realFilenames = readDirectory(realWordsPath);

  vector<Test> tests = vector<Test>();

  for(unsigned index = 0; index < realFilenames.size(); index++)
  {   
    vector<WordsLine> outputlines = readXml(ourFilenames.at(index));
    vector<WordsLine> truthlines  = readXml(realFilenames.at(index));
    Test test = Test(outputlines, truthlines, matchcontent);
    tests.push_back(test);
  }
  Test supertest = Test(tests);
  supertest.printResults();
  
  
}

void testNFold(int n, bool matchcontent){
  //TODO: read in files from both directories
  //      compare the output of our program with the correct output  
  cout << endl << "//// I'm going to test ////" << endl << endl;
  if(matchcontent){
    cout << "Matching content" << endl << endl;
  }else{
    cout << "Not matching content" << endl << endl;
  } 
  vector<Test> tests = vector<Test>();
  
  vector< string > realFilenames = readDirectory("../dataset/lines_train_all/", SHORT);
  for(size_t fileNr = 0; fileNr < realFilenames.size(); fileNr++){ 
    for(int nIndex = 0; nIndex < n; nIndex++){
      if(!isInNthTestSet(n, nIndex, fileNr)){
	continue;	    
      }      
      string trueFilePath = "../dataset/lines_train_all/" + realFilenames[fileNr];
      string outputFilePath = getNFoldOutputFolderName(n, nIndex) + "/" + realFilenames[fileNr];
      vector<WordsLine> outputlines = readXml(outputFilePath);
      vector<WordsLine> truthlines  = readXml(trueFilePath);
      Test test = Test(outputlines, truthlines, matchcontent);
      tests.push_back(test);
    }	
  }
  Test supertest = Test(tests);
  supertest.printResults();
  
  writeToResult(supertest.getResults());
}

void writeToResult(string result){
  ofstream resultFile("result.txt", ios::out | ios::app);
  resultFile << result;
  resultFile.close();
}

void runParameterLoop(){
    runAllNFold(10);
}

void runAllNFold(int n){
  bool allDatabasesExist = true;
  for(int nIndex = 0; nIndex < n; nIndex++){
    FILE * pFile;
    pFile = fopen (getNFoldDatabaseName(n,nIndex).c_str(),"r");    
    if (pFile == NULL) {
      allDatabasesExist = false;
      break;
    }else{
      fclose(pFile);
    }
  }
  
  char useFile = 'n';
  if (allDatabasesExist){
  //  cout << "The databases for "<< n <<"fold already exists." << endl;
  //  cout << "Do you want to skip training and use these databases? (warning: the databases will be overwritten if you choose 'n') (y/n) ";
  //  cin >> useFile;
    useFile = 'y';
  }  
  if(useFile == 'n')
  { 
    //first step: create database from trainingset
    trainNFold("../dataset/images_train_all/", "../dataset/lines_train_all/", 10);
  }
    //second step: classify testset using databse
  vector< string > testFileNames = readDirectory("../dataset/images_train_all/", SHORT);
  for(size_t fileNr = 0; fileNr < testFileNames.size(); fileNr++){ 
    for(int nIndex = 0; nIndex < n; nIndex++){
      if(!isInNthTestSet(n, nIndex, fileNr)){
	continue;	    
      }      
      string outputFolderName = getNFoldOutputFolderName(n, nIndex);
      mkdir(outputFolderName.c_str(),0777);
      string filename = testFileNames[fileNr].substr(0, testFileNames[n].length() - 4);          
      recognize(getNFoldDatabaseName(n,nIndex), "../dataset/images_train_all/" + filename + ".ppm", "../dataset/lines_train_all/" + filename + ".words", outputFolderName + "/" + filename + ".words");
    }	
  }
  
  //third step: test the output files against the desired output
  //testNFold(n, false);
  testNFold(n, true);
  
}

void setConfigFile(string configFileName){
  Config::getInstance()->initialize(configFileName);
  writeToResult("Loading config file: "+configFileName);
}

void runAll()
{
  FILE * pFile;
  pFile = fopen ("./database.db","r");
  char useFile = 'n';
  if (pFile != NULL)
  {
    cout << "A database file with the name database.db already exists." << endl;
    cout << "Do you want to skip training and use this file? (warning: database will be overwritten if you choose 'n') (y/n) ";
    cin >> useFile;
    fclose(pFile);
  }
  if(useFile == 'n')
  { 
    //first step: create database from trainingset
    train("../dataset/images_train/", "../dataset/lines_train/", "./database.db");
  }
 
  //second step: classify testset using databse
  vector< string > testFileNames = readDirectory("../dataset/images_test/", SHORT);
  for(size_t n = 0; n < testFileNames.size(); n++)
  { 
    string filename = testFileNames[n].substr(0, testFileNames[n].length() - 4);
    recognize("./database.db", "../dataset/images_test/" + filename + ".ppm", "../dataset/lines_test/" + filename + ".words", 
	      "./output/" + filename + ".words");
  }
  
  //third step: test the output files against the desired output
  test("../dataset/lines_test/", "./output/", false);
  test("../dataset/lines_test/", "./output/", true);
}

void saveDatabaseFeaturesAsCSV(){
  writeAllToCSV(readDatabase("./database.db"));
}

void saveDatabaseFeatureHistogramsAsCSV(){
  writeAllHistogramsToCSV(readDatabase("./database.db"),100);
}
